// Core Extension Implementation - 核心扩展实现
// 基于 augment-code/out/extension.js 的逆向工程和 libs/tools 的集成

import { createAnthropic } from '@augment/chat-model';
import { ModelProvider, SendRequest, StructuredNode, NodeType } from '@augment/chat-core';
import { ToolManager, ToolExecutionContext, ToolResult, ToolCategory } from '@augment/tools';
import Anthropic from '@anthropic-ai/sdk';

/**
 * 核心扩展类 - 整合所有工具和业务逻辑
 */
export class AugmentCoreExtension implements ModelProvider {
    private anthropic: Anthropic;
    private toolManager: ToolManager;
    private cancelledRequests = new Set<string>();
    private logger: any;
    private initialized: boolean = false;

    constructor(
        apiKey: string, 
        private modelName: string = 'claude-3-5-sonnet-20241022',
        logger?: any
    ) {
        this.logger = logger || console;
        this.anthropic = createAnthropic('anthropic', apiKey);
        this.toolManager = new ToolManager(this.logger);
    }

    /**
     * 初始化扩展
     */
    async initialize(): Promise<void> {
        if (this.initialized) {
            this.logger.warn('AugmentCoreExtension is already initialized');
            return;
        }

        try {
            this.logger.info('Initializing AugmentCoreExtension...');

            // 初始化工具系统
            await this.toolManager.initialize();

            this.initialized = true;
            this.logger.info(`AugmentCoreExtension initialized successfully with tools:`, 
                this.toolManager.getAvailableTools());

        } catch (error) {
            this.logger.error('Failed to initialize AugmentCoreExtension:', error);
            throw error;
        }
    }

    /**
     * 获取可用工具列表（Anthropic API 格式）
     */
    private getAnthropicTools(): any[] {
        const availableTools = this.toolManager.getAvailableTools();
        const anthropicTools: any[] = [];

        // 文件操作工具
        if (availableTools.includes('read-file')) {
            anthropicTools.push({
                name: 'read-file',
                description: 'Read content from a file',
                input_schema: {
                    type: 'object',
                    properties: {
                        path: { type: 'string', description: 'File path to read' },
                        encoding: { type: 'string', description: 'File encoding (default: utf8)' }
                    },
                    required: ['path']
                }
            });
        }

        if (availableTools.includes('save-file')) {
            anthropicTools.push({
                name: 'save-file',
                description: 'Save content to a file',
                input_schema: {
                    type: 'object',
                    properties: {
                        path: { type: 'string', description: 'File path to save' },
                        content: { type: 'string', description: 'Content to save' },
                        encoding: { type: 'string', description: 'File encoding (default: utf8)' },
                        add_last_line_newline: { type: 'boolean', description: 'Add newline at end (default: true)' }
                    },
                    required: ['path', 'content']
                }
            });
        }

        if (availableTools.includes('str-replace-editor')) {
            anthropicTools.push({
                name: 'str-replace-editor',
                description: 'Edit file by replacing specific text',
                input_schema: {
                    type: 'object',
                    properties: {
                        path: { type: 'string', description: 'File path to edit' },
                        old_str: { type: 'string', description: 'Text to replace' },
                        new_str: { type: 'string', description: 'Replacement text' },
                        old_str_start_line_number: { type: 'number', description: 'Start line number' },
                        old_str_end_line_number: { type: 'number', description: 'End line number' }
                    },
                    required: ['path', 'old_str', 'new_str', 'old_str_start_line_number', 'old_str_end_line_number']
                }
            });
        }

        // 进程管理工具
        if (availableTools.includes('launch-process')) {
            anthropicTools.push({
                name: 'launch-process',
                description: 'Launch a system process',
                input_schema: {
                    type: 'object',
                    properties: {
                        command: { type: 'string', description: 'Command to execute' },
                        wait: { type: 'boolean', description: 'Whether to wait for completion' },
                        max_wait_seconds: { type: 'number', description: 'Maximum wait time in seconds' },
                        cwd: { type: 'string', description: 'Working directory' }
                    },
                    required: ['command', 'wait', 'max_wait_seconds', 'cwd']
                }
            });
        }

        if (availableTools.includes('kill-process')) {
            anthropicTools.push({
                name: 'kill-process',
                description: 'Kill a running process',
                input_schema: {
                    type: 'object',
                    properties: {
                        terminal_id: { type: 'number', description: 'Terminal ID of the process to kill' }
                    },
                    required: ['terminal_id']
                }
            });
        }

        if (availableTools.includes('read-process')) {
            anthropicTools.push({
                name: 'read-process',
                description: 'Read output from a process',
                input_schema: {
                    type: 'object',
                    properties: {
                        terminal_id: { type: 'number', description: 'Terminal ID of the process' },
                        wait: { type: 'boolean', description: 'Whether to wait for completion' },
                        max_wait_seconds: { type: 'number', description: 'Maximum wait time in seconds' }
                    },
                    required: ['terminal_id', 'wait', 'max_wait_seconds']
                }
            });
        }

        if (availableTools.includes('write-process')) {
            anthropicTools.push({
                name: 'write-process',
                description: 'Write input to a process',
                input_schema: {
                    type: 'object',
                    properties: {
                        terminal_id: { type: 'number', description: 'Terminal ID of the process' },
                        input_text: { type: 'string', description: 'Text to write to the process' }
                    },
                    required: ['terminal_id', 'input_text']
                }
            });
        }

        if (availableTools.includes('list-processes')) {
            anthropicTools.push({
                name: 'list-processes',
                description: 'List all running processes',
                input_schema: {
                    type: 'object',
                    properties: {},
                    required: []
                }
            });
        }

        // 记忆管理工具
        if (availableTools.includes('remember')) {
            anthropicTools.push({
                name: 'remember',
                description: 'Create a memory item',
                input_schema: {
                    type: 'object',
                    properties: {
                        memory: { type: 'string', description: 'Content to remember' }
                    },
                    required: ['memory']
                }
            });
        }

        if (availableTools.includes('memory-retrieval')) {
            anthropicTools.push({
                name: 'memory-retrieval',
                description: 'Retrieve memories based on query',
                input_schema: {
                    type: 'object',
                    properties: {
                        query: { type: 'string', description: 'Search query for memories' },
                        limit: { type: 'number', description: 'Maximum number of results (default: 10)' }
                    },
                    required: ['query']
                }
            });
        }

        return anthropicTools;
    }

    /**
     * 执行工具调用
     */
    private async executeToolCall(toolName: string, input: any): Promise<ToolResult> {
        const context: ToolExecutionContext = {
            workspaceRoot: process.cwd(),
            environment: process.env as Record<string, string>,
            timestamp: new Date(),
            permissions: this.getPermissionsForTool(toolName),
            logger: this.logger
        };

        try {
            // 预处理参数
            const processedInput = this.preprocessToolInput(toolName, input, context);
            
            // 执行工具
            const result = await this.toolManager.executeTool(toolName, processedInput, context);
            
            this.logger.info(`Tool ${toolName} executed successfully`, { 
                success: result.success, 
                hasData: !!result.data 
            });
            
            return result;

        } catch (error) {
            this.logger.error(`Tool ${toolName} execution failed:`, error);
            return {
                success: false,
                error: {
                    code: 'EXECUTION_FAILED',
                    message: error instanceof Error ? error.message : `${error}`
                }
            };
        }
    }

    /**
     * 预处理工具输入参数
     */
    private preprocessToolInput(toolName: string, input: any, context: ToolExecutionContext): any {
        const processedInput = { ...input };

        // 为 launch-process 工具添加默认参数
        if (toolName === 'launch-process') {
            processedInput.wait = processedInput.wait !== undefined ? processedInput.wait : true;
            processedInput.max_wait_seconds = processedInput.max_wait_seconds !== undefined ? processedInput.max_wait_seconds : 30;
            processedInput.cwd = processedInput.cwd !== undefined ? processedInput.cwd : context.workspaceRoot;
        }

        // 为 memory-retrieval 工具添加默认参数
        if (toolName === 'memory-retrieval') {
            processedInput.limit = processedInput.limit !== undefined ? processedInput.limit : 10;
        }

        return processedInput;
    }

    /**
     * 获取工具权限
     */
    private getPermissionsForTool(toolName: string): string[] {
        const permissions: string[] = [];

        // 文件操作权限
        if (toolName.includes('file') || ['read-file', 'save-file', 'str-replace-editor'].includes(toolName)) {
            permissions.push('file:read', 'file:write');
        }

        // 进程管理权限
        if (toolName.includes('process') || ['launch-process', 'kill-process', 'read-process', 'write-process', 'list-processes'].includes(toolName)) {
            permissions.push('process:execute', 'process:read', 'process:write', 'process:kill', 'process:list');
        }

        // 记忆管理权限
        if (['remember', 'memory-retrieval'].includes(toolName)) {
            permissions.push('memory:read', 'memory:write');
        }

        return permissions;
    }

    /**
     * 格式化工具执行结果
     */
    private formatToolResult(toolName: string, result: ToolResult): string {
        if (!result.success) {
            return `❌ 工具 ${toolName} 执行失败: ${result.error?.message || '未知错误'}`;
        }

        switch (toolName) {
            case 'read-file':
                if (result.data?.content) {
                    const content = result.data.content;
                    const preview = content.length > 500 ? content.substring(0, 500) + '...' : content;
                    return `📄 文件读取成功 (${result.data.size || content.length} 字符):\n\`\`\`\n${preview}\n\`\`\``;
                }
                return `✅ 文件读取成功`;

            case 'save-file':
                return `✅ 文件保存成功: ${result.data?.path || ''}`;

            case 'str-replace-editor':
                return `✅ 文件编辑成功: 替换了 ${result.data?.linesReplaced || 0} 行`;

            case 'launch-process':
                if (result.data?.terminal_id) {
                    const output = result.data.output ? `\n输出:\n\`\`\`\n${result.data.output}\n\`\`\`` : '';
                    const error = result.data.error ? `\n错误:\n\`\`\`\n${result.data.error}\n\`\`\`` : '';
                    return `🚀 进程启动成功 (终端ID: ${result.data.terminal_id})${output}${error}`;
                }
                return `✅ 进程执行成功`;

            case 'kill-process':
                return `🛑 进程终止成功: ${result.data?.message || ''}`;

            case 'read-process':
                const output = result.data?.output ? `\n输出:\n\`\`\`\n${result.data.output}\n\`\`\`` : '';
                const error = result.data?.error ? `\n错误:\n\`\`\`\n${result.data.error}\n\`\`\`` : '';
                return `📖 进程输出读取成功${output}${error}`;

            case 'write-process':
                return `✍️ 进程输入写入成功: ${result.data?.message || ''}`;

            case 'list-processes':
                const processes = result.data?.processes || [];
                const processInfo = processes.map((p: any) =>
                    `- 终端ID: ${p.terminal_id}, 状态: ${p.completed ? '已完成' : '运行中'}`
                ).join('\n');
                return `📋 进程列表 (${processes.length} 个进程):\n${processInfo}`;

            case 'remember':
                return `💾 记忆创建成功: ${result.data?.message || ''}`;

            case 'memory-retrieval':
                const memories = result.data?.results || [];
                const memoryInfo = memories.slice(0, 3).map((m: any) =>
                    `- ${m.content.substring(0, 100)}${m.content.length > 100 ? '...' : ''}`
                ).join('\n');
                const moreInfo = memories.length > 3 ? `\n... 还有 ${memories.length - 3} 条记忆` : '';
                return `🧠 找到 ${memories.length} 条相关记忆:\n${memoryInfo}${moreInfo}`;

            default:
                return `✅ 工具 ${toolName} 执行成功`;
        }
    }

    /**
     * 生成响应 - 实现 ModelProvider 接口
     */
    async *generate(req: SendRequest & { requestId: string }): AsyncIterable<StructuredNode> {
        const { requestMessage, requestId } = req;

        if (!this.initialized) {
            yield {
                type: NodeType.RawResponse,
                ts: Date.now(),
                requestId,
                text: '❌ 扩展未初始化，请先调用 initialize() 方法'
            };
            return;
        }

        try {
            // 发送思考节点
            yield {
                type: NodeType.Thinking,
                ts: Date.now(),
                requestId,
                text: 'Analyzing request and preparing tools...'
            };

            // 检查是否被取消
            if (this.cancelledRequests.has(requestId)) {
                return;
            }

            // 调用 Anthropic API
            const response = await this.anthropic.messages.create({
                max_tokens: 4096,
                messages: [
                    { role: 'user', content: requestMessage }
                ],
                model: this.modelName,
                temperature: 0.7,
                tools: this.getAnthropicTools()
            });

            // 处理响应内容
            for (const content of response.content) {
                if (this.cancelledRequests.has(requestId)) {
                    return;
                }

                if (content.type === 'text') {
                    // 发送文本响应节点
                    yield {
                        type: NodeType.RawResponse,
                        ts: Date.now(),
                        requestId,
                        text: content.text
                    };
                } else if (content.type === 'tool_use') {
                    // 发送工具使用节点
                    yield {
                        type: NodeType.ToolUse,
                        ts: Date.now(),
                        requestId,
                        tool: content.name,
                        input: content.input
                    };

                    // 执行工具并返回结果
                    try {
                        const toolResult = await this.executeToolCall(content.name, content.input);
                        const resultText = this.formatToolResult(content.name, toolResult);

                        yield {
                            type: NodeType.RawResponse,
                            ts: Date.now(),
                            requestId,
                            text: resultText
                        };

                    } catch (error) {
                        yield {
                            type: NodeType.RawResponse,
                            ts: Date.now(),
                            requestId,
                            text: `工具执行异常: ${error instanceof Error ? error.message : '未知错误'}`
                        };
                    }
                }
            }

            // 生成建议问题
            yield {
                type: NodeType.SuggestedQuestions,
                ts: Date.now(),
                requestId,
                questions: ['继续解释', '举个例子', '还有其他方法吗？', '查看相关记忆']
            };

            // 标记主文本完成
            yield {
                type: NodeType.MainTextFinished,
                ts: Date.now(),
                requestId
            };

        } catch (error) {
            this.logger.error('Generation error:', error);
            yield {
                type: NodeType.RawResponse,
                ts: Date.now(),
                requestId,
                text: `错误: ${error instanceof Error ? error.message : String(error)}`
            };

            yield {
                type: NodeType.MainTextFinished,
                ts: Date.now(),
                requestId
            };
        } finally {
            this.cancelledRequests.delete(requestId);
        }
    }

    /**
     * 取消请求
     */
    cancel(requestId: string): void {
        this.cancelledRequests.add(requestId);
        this.logger.info(`Request ${requestId} cancelled`);
    }

    /**
     * 获取工具统计信息
     */
    getToolStats(): any {
        return this.toolManager.getToolStats();
    }

    /**
     * 启用/禁用工具
     */
    setToolEnabled(toolName: string, enabled: boolean): boolean {
        return this.toolManager.setToolEnabled(toolName, enabled);
    }

    /**
     * 检查工具是否可用
     */
    isToolAvailable(toolName: string): boolean {
        return this.toolManager.isToolAvailable(toolName);
    }

    /**
     * 获取可用工具列表
     */
    getAvailableTools(): string[] {
        return this.toolManager.getAvailableTools();
    }

    /**
     * 按类别获取工具
     */
    getToolsByCategory(category: ToolCategory): string[] {
        return this.toolManager.getToolsByCategory(category);
    }

    /**
     * 清理资源
     */
    async dispose(): Promise<void> {
        try {
            this.logger.info('Disposing AugmentCoreExtension...');
            await this.toolManager.cleanup();
            this.cancelledRequests.clear();
            this.initialized = false;
            this.logger.info('AugmentCoreExtension disposed successfully');
        } catch (error) {
            this.logger.error('Failed to dispose AugmentCoreExtension:', error);
        }
    }
}
